#!/usr/bin/python3
# coding=utf-8
#
# Copyright (C) 2023-2024. Huawei Technologies Co., Ltd. All rights reserved.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# ===============================================================================

import numpy as np

def linspace_custom(start, end, num, output_dtype=np.float32):
    if num <= 1:
        return np.array([start], dtype=output_dtype)
    start_f32 = start.astype(np.float32)
    end_f32 = end.astype(np.float32)
    step = (end_f32 - start_f32) / (num - 1)
    seq = start_f32 + np.arange(num, dtype=np.float32) * step
    return seq.astype(output_dtype)

def gen_golden_data_linspace():
    output_dtype = np.float32   
    ns_dtype = np.int64 
    # ==================== 固定数据类型配置 ====================
    # input_dtype1 = np.float16    
    input_dtype1 = np.float32
    # input_dtype1 = np.int8
    # input_dtype1 = np.int16
    # input_dtype1 = np.int32
    # input_dtype2 = np.float16    
    input_dtype2 = np.uint8
    # input_dtype2 = np.int8
    # input_dtype2 = np.int16
    # input_dtype2 = np.int32
    # =======================基础参数配置======================
    min_n = 1                         # n的最小值
    max_n = 1024                     # n的最大值
    # start_range = (-1000.0, 1000.0)   # start的取值范围
    start_range = (0, 1000.0)
    end_offset_range = (0, 255)  # end与start的差值范围
    # ========================================================
    starts = np.random.uniform(*start_range, size=1).astype(input_dtype1)
    end_offset = np.random.uniform(*end_offset_range, size=1).astype(input_dtype2)
    ends = (starts + end_offset).astype(input_dtype2)
    ns = np.random.randint(min_n, max_n + 1, size=1, dtype=ns_dtype)
    # golden = np.linspace(starts[0], ends[0], ns[0]).astype(output_dtype)
    # 注意np库中的linspace输出数据类型为double，强转为float32会损失精度
    golden = linspace_custom(starts[0], ends[0], ns[0], output_dtype=output_dtype)
    starts.tofile("./input/starts.bin")
    ends.tofile("./input/ends.bin")
    ns.tofile("./input/ns.bin")
    golden.tofile("./output/golden.bin")
    
if __name__ == "__main__":
    gen_golden_data_linspace()